'''
CT 图像HU值转换
CT 图像标准化
CTSlice 生成
补充CTSlice文件地址到csv文件
'''

import pandas
import pydicom
import numpy as np
import csv
import os

# 将输入图像的像素值（-1024，2000）归一化到0-1之间
def normalize_hu(image):
    # MIN_BOUND = -1000.0
    # MAX_BOUND = 400.0
    MIN_BOUND = -1350.0
    MAX_BOUND = 150.0
    image = (image - MIN_BOUND) / (MAX_BOUND - MIN_BOUND)
    image[image > 1] = 1.
    image[image < 0] = 0.
    return image

# 提取CT图像素值（-4000，4000），CT图的像素值是由HU值表示的
def get_pixels_hu(slice, intercept, slope):
    image = slice.pixel_array
    image = image.astype(np.int16)
    image[image == -2000] = 0
    if slope != 1:
        image = slope * image.astype(np.float64)
        image = image.astype(np.int16)
    image += np.int16(intercept)
    return np.array(image, dtype=np.int16)

def generate_CTSlice():
    origin_path = 'D:/lung_cancer/data/data_augmentation/all_data2.csv'
    data = pandas.read_csv(origin_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        instancenumber = str(data['z'][i])
        ct_name = str(data['CT_origin_path'][i]).split('/')[-1]

        ct_path = 'D:/lung_cancer/data/data_augmentation/origin_data/' + str(patient) + '/CT/' + ct_name

        ct_intercept = data['ct_intercept'][i]
        ct_slope = data['ct_slope'][i]
        cancer_type = data['cancer_type'][i]


        slice = pydicom.read_file(ct_path)
        ct_array = get_pixels_hu(slice, ct_intercept, ct_slope)
        normalized_array = normalize_hu(ct_array)


        save_path = 'D:/lung_cancer/data/data_augmentation/Slice/'+str(patient)+'/CTSlice/'
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        np.save(save_path+instancenumber+'.npy', normalized_array)
        print('%d--->patient: %s' % (i+1, str(patient)))


# 将生成的ct_array补充到all_data.csv中
def add_ctpath():
    data_path = 'D:/lung_cancer/data/data_augmentation/all_data2.csv'
    data = []
    f = csv.reader(open(data_path, 'r'))
    for i in f:
        data.append(i)

    new_data = []

    for line in data[1:]:
        patient = str(line[0])
        cancer_type = str(line[5])
        instancenumber = str(line[1])
        ct_path = 'Slice/'+patient+'/CTSlice/'+instancenumber+'.npy'
        line.append(ct_path)
        new_data.append(line)
    data[0].append('CTSlice_Path')
    df = pandas.DataFrame(new_data, columns=data[0])
    df.to_csv('D:/lung_cancer/data/data_augmentation/all_data3.csv', index=False)


if __name__ == '__main__':
    # generate_CTSlice()
    add_ctpath()